import ptan
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

HID_SIZE = 128


class DDPGActor(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGActor, self).__init__()

        # 构建网络
        if isinstance(obs_size, (np.ndarray, list, tuple)):
            input_dim = np.prod(obs_size)

        self.net = nn.Sequential(
            nn.Linear(input_dim, 400),
            nn.ReLU(),
            nn.Linear(400, 300),
            nn.ReLU(),
            nn.Linear(300, act_size),
            nn.Tanh()
        )

        self._init_weights()

    def _init_weights(self):
        # 小范围初始化最后一层，防止初始输出过大
        nn.init.uniform_(self.net[-2].weight, -3e-3, 3e-3)
        nn.init.uniform_(self.net[-2].bias, -3e-3, 3e-3)


    def forward(self, x):
        x = x.view(x.size(0), -1)
        return self.net(x)


class DDPGCritic(nn.Module):
    '''
    深度确定性策略梯度网络Q值评价网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGCritic, self).__init__()

        if isinstance(obs_size, (np.ndarray, list, tuple)):
            input_dim = np.prod(obs_size)

        # 构建环境观察网络
        self.obs_net = nn.Sequential(
            nn.Linear(input_dim, 400),
            nn.ReLU(),
        )

        # 构建观察+动作合并网路，输出评价Q值
        # 以前的Q值网络是输出每个动作的Q值，在连续值里面是直接输出评价Q值
        # 这里的400是self.obs_net的输出维度，act_size是动作网络的输出维度
        # 为了后续合并预测评价做准备
        self.out_net = nn.Sequential(
            nn.Linear(400 + act_size, 300),
            nn.ReLU(),
            nn.Linear(300, 1)
        )

    def forward(self, x, a):
        x = x.view(x.size(0), -1)
        a = a.view(a.size(0), -1)
        obs = self.obs_net(x)
        return self.out_net(torch.cat([obs, a], dim=1))



class AgentDDPG(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, net, device="cpu", ou_enabled=True, ou_mu=0.0, ou_teta=0.15, ou_sigma=0.2, ou_epsilon=1.0):
        self.net = net
        self.device = device
        self.ou_enabled = ou_enabled
        self.ou_mu = ou_mu
        self.ou_teta = ou_teta
        self.ou_sigma = ou_sigma
        self.ou_epsilon = ou_epsilon
        self.agent_states = None

    def initial_state(self):
        # 统一接口，但是这里不需要使用，返回None而不是使用pass
        return None

    def __call__(self, states, agent_states):
        '''
        states：当前的环境状态
        agent_states: 内部智能体的状态,之前的代理器里面这个基本无用,一开始的时候，agent_states是空的，但在这里因为要使用OU过程对执行的动作进行噪音干扰，所以需要使用了智能体的内部状态
        '''
        # 将环境转换为目标的数据类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 得到执行的动作输出
        mu_v = self.net(states_v)
        actions = mu_v.data.cpu().numpy()

        if self.ou_enabled and self.ou_epsilon > 0:
            # 启动随机探索
            new_a_states = []
            for a_state, action in zip(agent_states, actions):
                if a_state is None:
                    # 为什么智能体状态为空时可以构建一个维度和动作相同，且全为零的状态
                    # 这是针对初始状态，初始情况下，智能体的状态为空，则创建一个为0的状态
                    a_state = np.zeros(shape=action.shape, dtype=np.float32)
                # 书p298也 todo 了解OU过程
                a_state += self.ou_teta * (self.ou_mu - a_state)
                a_state += self.ou_sigma * np.random.normal(size=action.shape)

                action += self.ou_epsilon * a_state
                new_a_states.append(a_state)
        else:
            new_a_states = agent_states

        # 这个步骤 的意思是修正动作值到-1和1之间，否则将导致PyBullet抛出异常
        actions = np.clip(actions, -1, 1)
        self.agent_states = new_a_states
        return actions.reshape(actions.shape[0], 1, actions.shape[1]), new_a_states

    def state_dict(self):
        return {
            'ou_enabled': self.ou_enabled,
            'ou_mu': self.ou_mu,
            'ou_teta': self.ou_teta,
            'ou_sigma': self.ou_sigma,
            'ou_epsilon': self.ou_epsilon,
            'agent_states': self.agent_states
        }

    def load_state_dict(self, state_dict):
        self.ou_enabled = state_dict['ou_enabled']
        self.ou_mu = state_dict['ou_mu']
        self.ou_teta = state_dict['ou_teta']
        self.ou_sigma = state_dict['ou_sigma']
        self.ou_epsilon = state_dict['ou_epsilon']
        self.agent_states = state_dict['agent_states']


class ModelActor(nn.Module):
    def __init__(self, obs_size, act_size):
        '''
        :param obs_size: 观测的环境维度
        :param act_size: 执行的动作的维度
        '''
        super(ModelActor, self).__init__()

        self.mu = nn.Sequential(
            nn.Linear(obs_size, HID_SIZE),
            nn.Tanh(),
            nn.Linear(HID_SIZE, HID_SIZE),
            nn.Tanh(),
            nn.Linear(HID_SIZE, act_size),
            nn.Tanh(),
        )
        # 作用 看名称像是方差：是用来控制动作的探索程度的
        # 怎么更新？：在训练的过程中，会不断的更新这个参数，更新的逻辑就在于计算熵损失以及计算动作优势大小的时候会参与计算，然后在梯度更新的时候，会自动更新这个参数到合适的大小
        self.logstd = nn.Parameter(torch.zeros(act_size))

    def forward(self, x):
        return self.mu(x)


class ModelCritic(nn.Module):
    '''
    trop信赖域策略优化评价网络
    ACKTR算法中使用的critic网络
    ppt优化评价网络
    '''
    def __init__(self, obs_size):
        super(ModelCritic, self).__init__()

        self.value = nn.Sequential(
            nn.Linear(obs_size, HID_SIZE),
            nn.ReLU(),
            nn.Linear(HID_SIZE, HID_SIZE),
            nn.ReLU(),
            nn.Linear(HID_SIZE, 1),
        )

    def forward(self, x):
        return self.value(x)

class AgentAcktr(ptan.agent.BaseAgent):
    '''
    创建代理器
    '''
    def __init__(self, net, device="cpu"):
        self.net = net
        self.device = device

    def __call__(self, states, agent_states):
        '''
        states: 观测的环境状态
        agent_states：智能体自己的状态，在这里是没有使用的
        '''
        # 创建环境预处理器，将环境状态转换为float32类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)

        # 通过环境状态预测执行的动作
        mu_v = self.net(states_v)
        mu = mu_v.data.cpu().numpy()
        logstd = self.net.logstd.data.cpu().numpy()
        # 该动作的作用，是对预测的动作添加随机噪音，实现动作的探索
        actions = mu + np.exp(logstd) * np.random.normal(size=logstd.shape)
        # 将执行的动作压缩到-1到1中，可能是因为输入给网络的值不能超过-1和1
        actions = np.clip(actions, -1, 1)
        return actions, agent_states
    

class LinearFeatureExtractor(nn.Module):

    def __init__(self, obs_size):
        super().__init__()

        if isinstance(obs_size, (np.ndarray, list, tuple)):
            input_dim = np.prod(obs_size)

        self.net = nn.Sequential(
            nn.Linear(input_dim, input_dim * 2),
            nn.BatchNorm1d(input_dim * 2),
            nn.LeakyReLU(),
            nn.Linear(input_dim * 2, input_dim * 2),
            nn.BatchNorm1d(input_dim * 2),
            nn.LeakyReLU()
        )

        self.out_size = input_dim * 2

    def forward(self, x):
        x = x.view(x.size(0), -1)
        return self.net(x)


class AcktrActor(nn.Module):
    def __init__(self, obs_size, act_size):
        '''
        :param obs_size: 观测的环境维度
        :param act_size: 执行的动作的维度
        '''
        super(AcktrActor, self).__init__()

        self.feature_extractor = LinearFeatureExtractor(obs_size)

        self.act = []
        for _ in range(act_size[0]):
            self.act.append(nn.Sequential(
                nn.Linear(self.feature_extractor.out_size, HID_SIZE),
                nn.LeakyReLU(),
                nn.Linear(HID_SIZE, act_size[1]),
                nn.Tanh(),
            ))
        # 作用 看名称像是方差：是用来控制动作的探索程度的
        # 怎么更新？：在训练的过程中，会不断的更新这个参数，更新的逻辑就在于计算熵损失以及计算动作优势大小的时候会参与计算，然后在梯度更新的时候，会自动更新这个参数到合适的大小
        self.logstd = nn.Parameter(torch.zeros(act_size))

    def forward(self, x):
        return torch.stack([a(self.feature_extractor(x)) for a in self.act], dim=1)


class AcktrCritic(nn.Module):
    '''
    trop信赖域策略优化评价网络
    ACKTR算法中使用的critic网络
    ppt优化评价网络
    '''
    def __init__(self, obs_size):
        super(AcktrCritic, self).__init__()

        self.feature_extractor = LinearFeatureExtractor(obs_size)

        self.value = nn.Sequential(
            nn.Linear(self.feature_extractor.out_size, HID_SIZE),
            nn.ReLU(),
            nn.Linear(HID_SIZE, 1),
        )

    def forward(self, x):
        return self.value(self.feature_extractor(x))